import torch.nn as nn
from monai.losses import DiceCELoss


class MyDiceCELoss(nn.Module):
    def __init__(self, to_onehot_y=False, softmax=False):
        super().__init__()

        self.loss = DiceCELoss(to_onehot_y=to_onehot_y, softmax=softmax)


    def forward(self, pred_stage1, pred_stage2, target):
        """
        :param pred_stage1: (B, 14, 48, 256, 256)
        :param pred_stage2: (B, 14, 48, 256, 256)
        :param target: (B, 48, 256, 256)
        """

        # Caculate cross entropy loss
        loss_stage1 = self.loss(pred_stage1, target)
        loss_stage2 = self.loss(pred_stage2, target)

        # The eventual loss is constructed by two parts of loss.
        return (loss_stage1 + loss_stage2) / 2
